# coding=utf-8
import os
import shutil
import threading
from queue import Queue

import pandas as pd
from loguru import logger

from cal_ops.point import cal_point1_txt
from mylib import download_all
from mylib.mydb import MyDB

st_queue = Queue()


def get_all_txt_files():
    for root, dirs, files in os.walk('txts'):
        return [(os.path.join(root, item), item.replace('.txt', '')) for item in files]


def clear_dir(d='result'):
    if os.path.exists(d):
        shutil.rmtree(d)
    os.makedirs(d)


def get_all_stock_csv_path():
    for root, dirs, files in os.walk('stocks'):
        return [os.path.join(root, item) for item in files]


def get_all_stocks_indus():
    with open('industry_count_rank.csv', 'r') as fr:
        all_indus = [item.strip().split(',') for item in fr.readlines()[1:]]
        return all_indus


def get_avg_top(indus, stocks, stocks_dict, N=50):
    stocks_top_dict = {}
    for stock in stocks:
        csv_path = f'stocks/{stock}.csv'
        if not os.path.exists(csv_path):
            continue
        df2 = pd.read_csv(csv_path)
        stock_k = f'{stock}_{stocks_dict.get(stock)}'
        stocks_top_dict[stock_k] = round(abs(df2[0:N]['pct_chg'].mean()), 4)
    sorted_d = sorted(stocks_top_dict.items(), key=lambda x: x[1], reverse=True)
    logger.error(f'{indus}近{N}日振幅排名')
    for item in sorted_d:
        logger.error(item)


def get_end_date(log_dir, indus, indus_num):
    f_path = f'{log_dir}/{indus_num}_all_{indus}.csv'
    if not os.path.exists(f_path):
        return None
    with open(f_path, 'r') as fr:
        line = fr.readline()
        return eval(line.split(',')[0])


def t_download():
    txts = get_all_txt_files()
    mydb_obj = MyDB()
    for txt_path, txt_name in txts:
        with open(txt_path) as fr:
            stocks_names = [item.strip() for item in fr.readlines() if item.strip()]
        if not stocks_names:
            continue
        try:
            stocks_and_names = mydb_obj.select_all_stocks_and_name_by_name(stocks_names)
        except Exception as e:
            logger.error(e)
            continue
        logger.warning(f't_download {txt_name} {stocks_and_names}')
        stocks = [i[0] for i in stocks_and_names]
        stocks = download_all.run(stocks)
        stocks_dict = {}
        for i in stocks_and_names:
            stocks_dict[i[0]] = i[1]
        st_queue.put((txt_name, len(stocks), stocks, stocks_dict))
    st_queue.put((1, 1, 1, 1))


def t_cal():
    while 1:
        if not st_queue:
            continue
        indus, indus_num, stocks, stocks_dict = st_queue.get()
        if (1, 1, 1, 1) == (indus, indus_num, stocks, stocks_dict):
            logger.warning('所有计算完成')
            break
        logger.warning(f't_cal get_avg_top 正在计算【{indus}, {indus_num}】')
        get_avg_top(indus, stocks, stocks_dict)
        logger.warning(f't_cal get_avg_top 计算完成【{indus}, {indus_num}】')


if __name__ == '__main__':
    # 计算振幅排名
    logger.add("rank_results/main_cal_zf_rank_{time}.txt", format='{message}', level='ERROR')
    t1 = threading.Thread(target=t_download, args=())
    t2 = threading.Thread(target=t_cal, args=())
    t2.start()
    t1.start()
    t2.join()
    t1.join()
